import mmap, time, ast, struct

###
# init
###

response.title = 'Firewall WebUI'



###
# constants
###

protocol = {
        1  : 'ICMP',
        6  : 'TCP',
        17 : 'UDP',
    }
for i in range(256):
    if not protocol.get(i): protocol[i] = i


reason = {
        -1  : 'firewall inactive',
        -2  : 'prot not enforced',
        -3  : 'prot not supported',
        -4  : 'no matching rule',
        -9  : 'out of state',
        -10 : 'conn tab full',
        -11 : 'app out of state',
        -12 : 'app attack',
    }

modes = {
        'active':      0x01,
        'icmp':        0x02,
        'tcp':         0x04,
        'udp':         0x08,
        'conn_track':  0x80,
        'app_ctl':     0x100,
    }

action = {
        0: 'DROP',
        1: 'ACCEPT',
    }

hook = {
    i: s for i, s in enumerate(['PRE', 'IN', 'FW', 'OUT', 'POST'])
    }


    
structs = {
    'log': 'I 3B 2I 2H i I',
    'rule': '>4B 2H 2I',
    'conn': '2I 2H I B 3x',
}

    
table_colors = {
    'row_even': '#DBE9FF',
    'row_odd': '#D8DADF',
    'title': '#A4AECE',
    }
    

conn_t = {
    1:  'Syn sent',
    2:  'Syn ack',
    3:  'Estab',
    4:  'Closing',
} 
    

###
# helpers
###

def int2ip(n):
    return '{}.{}.{}.{}'.format(
        (n >> 24) & 0xFF, (n >> 16) & 0xFF,
        (n >> 8) & 0xFF, n & 0xFF)
        
def ip2int(s):
    four_bytes = struct.pack('4B', *[int(i) for i in s.split('.')])
    return struct.unpack('>I', four_bytes)[0]
    
                               
###
# configuration
###


class Enforce:
    def enforcement(self, opt = None):
        with open('/sys/class/fw4/fw4_rules/enforcement') as f:
            e = ast.literal_eval(f.readline())
        if opt:
            e ^= modes[opt]
            with open('/sys/class/fw4/fw4_rules/enforcement', 'w') as f:
                f.write(str(e))
        
        return {k: (True if e & modes[k] else False) for k in modes.keys()}
                
    def make_enforce_button(self, opt):
        e = self.enforcement()
        return A(
            'ON' if e[opt] else 'OFF',
            _href = URL('conf', 'enforce', args = opt),
            cid = 'enforce_button_{}'.format(opt),
            _class = 'button',
            _style = 'color: {}'.format('green' if e[opt] else 'red'))
    
    


###
# log
###



class Log:
    def __init__(self, l):
        if not isinstance(l, str) or len(l) != struct.calcsize(structs['log']):
            raise Exception('bad log init: l = {}'.format(l))
        
        (self.time, self.prot, self.action, self.hooknum, 
            self.src_ip, self.dst_ip,                    
            self.src_port, self.dst_port,                
            self.reason, self.count) = struct.unpack(structs['log'], l)
        
        
    def make_row(self, i):
        
        return TR(TD(str(i)),
            TD(time.ctime(self.time)), TD(protocol[self.prot]), TD(action[self.action]),
            TD(hook[self.hooknum]), TD(int2ip(self.src_ip)), TD(str(self.src_port) if protocol[self.prot] != 'ICMP' else ''),
            TD(int2ip(self.dst_ip)), TD(str(self.dst_port) if protocol[self.prot] != 'ICMP' else ''), TD(str(self.count)),
            TD(reason.get(self.reason) or ('rule #' + str(self.reason))),
            _style = 'background-color: {}'.format(table_colors['row_even'] if i % 2 == 0 else table_colors['row_odd']))

class Logs:
    def __init__(self):
        self.logs = []
        with open('/dev/fw4_log', 'rb') as log, \
             open('/sys/class/fw4/fw4_log/log_size') as log_size_file:
     
             log_size = ast.literal_eval(log_size_file.readline())
             l = mmap.mmap(log.fileno(), log_size, mmap.MAP_SHARED, mmap.PROT_READ)
             
             i = 0
             while ord(l[i * 28 + 4]):
                 self.logs.append(Log(l[i * 28: (i + 1) * 28]))
                 i += 1

             l.close
        
    def get_logs_as_table(self, f, t):
        return TABLE(
            TR(TD('#'), TD('Time'), TD('Protocol'), TD('Action'), TD('Hook'), TD('Src IP'),
            TD('Src Port'), TD('Dest IP'), TD('Dest Port'), TD('Count'),
            TD('Reason'), _style = 'background-color: {}'.format(table_colors['title'])),
            *[l.make_row(i + f + 1) for i, l in enumerate(self.logs[f: t])])
    
    


###
# rules
###

class Rules:
    def __init__(self):
        self.rules = []
        with open('/dev/fw4_rules', 'r+b') as rule_base, \
            open('/sys/class/fw4/fw4_rules/rules_size') as rule_base_size_file:
            
            rule_base_size = ast.literal_eval(rule_base_size_file.readline())
            r = mmap.mmap(rule_base.fileno(), rule_base_size)
            i = 0
            while ord(r[i * 16]):
                self.rules.append(Rule(r[i * 16: (i + 1) * 16]))
                i += 1

            r.close()

    def get_as_table(self):
        return DIV(TABLE(
            TR(TD('#'), TD('Protocol'), TD('Src IP'), TD('Src Mask'), TD('Src Port'),
            TD('Dest IP'), TD('Dest Mask'), TD('Dest Port'), TD('Action'),
            TD('Edit', _id = 'rule_edit'), 
            _style = 'background-color: {}'.format(table_colors['title'])),
            *[r.make_row(i) for i, r in enumerate(self.rules)]),
            )
    
    def commit(self):
        with open('/dev/fw4_rules', 'r+b') as rule_base, \
            open('/sys/class/fw4/fw4_rules/rules_size') as rule_base_size_file:
            
            rule_base_size = ast.literal_eval(rule_base_size_file.readline())
            m = mmap.mmap(rule_base.fileno(), rule_base_size)
            
            for i, r in enumerate(self.rules + [Rule()]):
                m[i * 16: (i + 1) * 16] = r.make_bytes()
                        
            m.close()

class Rule:
    def __init__(self, r = chr(0) * 16, form = None):
        if not isinstance(r, str) or len(r) != struct.calcsize(structs['rule']):
            raise Exception('bad rule init, r = {}'.format(r))
        
        if form:
            self.protocol = int(form.vars.protocol)
            self.src_ip   = ip2int(form.vars.src_ip)
            self.src_mask = form.vars.src_mask
            self.src_port = form.vars.src_port
            self.dst_ip   = ip2int(form.vars.dst_ip)
            self.dst_mask = form.vars.dst_mask
            self.dst_port = form.vars.dst_port
            self.action   = int(form.vars.action)
        else:
            (self.protocol, self.src_mask, self.dst_mask, self.action,
                self.src_port, self.dst_port, self.src_ip,
                self.dst_ip) = struct.unpack(structs['rule'], r)

    def make_row(self, i):
        return TR(
            TD(str(i)), TD(protocol[self.protocol]), TD(int2ip(self.src_ip)), TD(str(self.src_mask)),
            TD(str(self.src_port) if ((protocol[self.protocol] in ['TCP', 'UDP']) and self.src_port != 0) else ''),
            TD(int2ip(self.dst_ip)), TD(str(self.dst_mask)),
            TD(str(self.dst_port) if ((protocol[self.protocol] in ['TCP', 'UDP']) and self.dst_port != 0) else ''),
            TD(str(action[self.action])),
            TD(self.make_crud(i), _id = 'rule_edit_{}'.format(i)),
            _style = 'background-color: {}'.format(table_colors['row_even'] if i % 2 == 0 else table_colors['row_odd']),
            _id = 'rule_{}'.format(i),
            )

    def make_bytes(self):
        return struct.pack(structs['rule'],
            self.protocol, self.src_mask, self.dst_mask, self.action,
            self.src_port, self.dst_port, self.src_ip, self.dst_ip)
             
         
    def make_crud(self, i):
        return DIV(*[A(s[:3], _href = URL('rule', s, args = i), cid = 'rules_display', _style = 'padding: 0px 5px') \
                     for s in ('delete', 'add', 'duplicate', 'edit', 'up', 'down')])
                
    def make_form(self):
        return SQLFORM.factory(
            Field('protocol', requires = IS_IN_SET(protocol, zero = None), default = self.protocol),
            Field('src_ip', requires = IS_IPV4(), default = int2ip(self.src_ip)),
            Field('src_mask', requires = IS_INT_IN_RANGE(0, 33), default = self.src_mask),
            Field('src_port', requires = IS_INT_IN_RANGE(0, 1 << 16), default = self.src_port),
            Field('dst_ip', requires = IS_IPV4(), default = int2ip(self.dst_ip)),
            Field('dst_mask', requires = IS_INT_IN_RANGE(0, 33), default = self.dst_mask),
            Field('dst_port', requires = IS_INT_IN_RANGE(0, 1 << 16), default = self.dst_port),
            Field('action', requires = IS_IN_SET(action, zero = None), default = self.action),
            )


###
# Connection table
###

"""// connection tracking
typedef struct {
    unsigned int cli_ip, ser_ip;
    unsigned short cli_port, ser_port;
    unsigned int expires;
    unsigned char state;
} connection_t;"""
